from limits.storage import MemoryStorage
from limits import RateLimitItemPerHour
from limits.strategies import FixedWindowRateLimiter


class RateLimit():
    """ IP或者APP请求限制 """

    per_hour = 60000
    storage: MemoryStorage
    strategy: FixedWindowRateLimiter
    hour_limit: RateLimitItemPerHour

    def __init__(self, per_hour=60000) -> None:
        """构造IP请求限制

        Args:
            per_hour (int, optional): 每个小时限制的请求数. Defaults to 60000.
        """
        self.per_hour = per_hour
        self.storage = MemoryStorage()  # 设置存储位置
        self.strategy = FixedWindowRateLimiter(self.storage)
        self.hour_limit = RateLimitItemPerHour(per_hour, 1)

    def hit(self, namespace):
        """ 判断是否被限制 """
        return self.strategy.hit(self.hour_limit, namespace)

    def close(self):
        self.storage = None
        self.strategy = None
        self.hour_limit = None


def limits_start(app):
    app.ctx.rate_limit = RateLimit(app.config['MAX_REQUEST_PER_HOUR'])


def limits_end(app):
    if app.ctx.rate_limit is not None:
        app.ctx.rate_limit.close()
        app.ctx.rate_limit = None
